#include "cuda_runtime.h"
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <curand.h>
#include "cudafunc.h"
#include "toroidal.h"
#include <cusparse.h>
#include <time.h>

extern field Apsi;
// complie command:
// nvcc curandtest.cu -lcurand -lcudart

__device__ fvec crotate(fvec a, float phi)
{return {a.x*cos(phi)+a.y*sin(phi),a.y*cos(phi)-a.x*sin(phi),a.z};}
__device__ float cdot(fvec a, fvec b)
{return a.x*b.x + a.y*b.y + a.z*b.z;}
__device__ fvec ccross(fvec a, fvec b)
{return {a.y*b.z-a.z*b.y, a.z*b.x-a.x*b.z, a.x*b.y-a.y*b.x};}
__device__ fvec ccyl2xyz(fcylv a)
{return {a.r*cos(a.p),a.r*sin(a.p),a.z};}
__device__ fcylv cxyz2cyl(fvec a)
{float pphi = atan(a.y/a.x) + pi*(a.x<0.0) + 2*pi*(a.y < 0.0 && a.x > 0.0);
return {sqrt(a.x*a.x + a.y*a.y), pphi, a.z};}
__device__ fvec cadd(fvec a, fvec b)
{return {a.x+b.x , a.y+b.y , a.z+b.z};}
__device__ fvec cmul(float a, fvec b)
{return {a*b.x, a*b.y, a*b.z};}

void cudarandgen(float* da, int n, char type)
{
    cudaDeviceSynchronize();
    curandGenerator_t gen;
    float* p_d;
    cudaMalloc((void**)&p_d, n*sizeof(float));
    curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_MRG32K3A);
    curandSetPseudoRandomGeneratorSeed(gen, 11ULL);
    if (type == 'u')
        curandGenerateUniform(gen, p_d, n);
    else if (type == 'g')
        curandGenerateNormal(gen, p_d, n, 0, 1);
    else
        return;
    cudaMemcpy(da, p_d, n*sizeof(float), cudaMemcpyDeviceToHost);
    curandDestroyGenerator(gen);
    cudaFree(p_d);
    
}

// interpolation in cylinder coordinate
__global__ void intpkernel(fpcrd* pc, float* A, float* pval, int numf, bool* ins)
{
    int idx = blockDim.x*blockIdx.x + threadIdx.x;
    if(ins[idx] == false)
        return;
    int rg = (int)((pc[idx].crd.r - Rmin)/dr);
    int zg = (int)((pc[idx].crd.z - Zmin)/dz);
    float ofr = (pc[idx].crd.r - (float)rg * dr - Rmin)/dr;
    float ofz = (pc[idx].crd.z - (float)zg * dz - Zmin)/dz;
    float coefrs[4] = {ofr*(ofr-(float)1.0)*(ofr-(float)2.0)/-(float)6.0,
     (ofr+(float)1.0)*(ofr-(float)1.0)*(ofr-(float)2.0)/(float)2.0,
    (ofr+(float)1.0)*ofr*(ofr-(float)2.0)/-(float)2.0,
     (ofr+(float)1.0)*ofr*(ofr-(float)1.0)/(float)6.0};
    float coefzs[4] = {ofz*(ofz-(float)1.0)*(ofz-(float)2.0)/-(float)6.0,
     (ofz+(float)1.0)*(ofz-(float)1.0)*(ofz-(float)2.0)/(float)2.0,
     (ofz+(float)1.0)*ofz*(ofz-(float)2.0)/-(float)2.0,
     (ofz+(float)1.0)*ofz*(ofz-(float)1.0)/(float)6.0};

    for (int k = 0; k < numf; k++)
        pval[idx*numf+k] = 0.0;
    for(int i = 0; i < 4; i++)
        for(int j = 0; j < 4; j++)
        {
            int point = 0;
            if(rg+i-1 < 0)
                point += Ny*(rg+i);
            else if(rg+i-1 > Nx)
                point += Ny*(rg+i-2);
            else
                point += Ny*(rg+i-1);
            if (zg+j-1 < 0)
                point += zg+j;
            else if (zg+j-1 > Nx)
                point += zg+j-2;
            else
                point += zg+j-1;
        for (int k = 0; k < numf; k++)
            pval[idx*numf+k] += A[Nx*Ny*k+point]*coefrs[i]*coefzs[j];
    }
}

__global__ void pushkernel(fpfvec* pfs, fphsp* pph, fpcrd* pcrdp, bool* ins, 
                     float dt, float q, float m)
{
    int id = blockDim.x*blockIdx.x + threadIdx.x;
    if(ins[id] == false)
        return;
    fvec E = pfs[id].E;
    fvec B = pfs[id].B;
    fvec gB = pfs[id].gB;
    fvec cb = pfs[id].cb;
    float vpa = pcrdp[id].vb;
    float mu = pcrdp[id].mu;
    fvec Bst = cadd(B , cmul(q*vpa/m,cb));
    float invBstn = sqrt(cdot(B,B))/cdot(Bst, B);    
    fvec dX = cmul(vpa,B);
    fvec curvdrift = cmul(q*vpa*vpa/m,cb);
    fvec graddrift = cmul(-pow(cdot(B,B),-0.5),ccross(B,cadd(E,cmul(-1*mu/q,gB))));
    dX = cadd(cadd(dX , curvdrift) , graddrift);
    dX = cmul(dt*invBstn , dX);
    float dv = cdot(Bst,cadd(E,cmul((-1*mu/q),gB)))*q*invBstn/m;
    pph[id] = {dX, dt*dv}; 
}

// rotate vector field and gather
__global__ void rogakernel(fpcrd* px, float* flds, fpfvec* pfs, feqdst* pFeq,
 bool h1, int n, float omega, float t)
{
    int id = blockDim.x*blockIdx.x + threadIdx.x;
    float angle = -px[id].crd.p;
    float temp[18];
    for (int i = 0; i < 18; i++)
        temp[i] = flds[id*18 + i];
    pfs[id].B = crotate({temp[0],temp[1],temp[2]},angle);
    pfs[id].gB = crotate({temp[3],temp[4],temp[5]},angle);
    pfs[id].cb = crotate({temp[6],temp[7],temp[8]},angle);
    if(h1)
    {
        fvec Ei = crotate({temp[9],temp[10],temp[11]},angle);
        fvec Eq = crotate({temp[12],temp[13],temp[14]},angle);
        float phase = -omega*t - angle*n;
        pfs[id].E = cadd(cmul(cos(phase),Ei), cmul(sin(phase),Eq));
    }
    else
        pfs[id].E = {0.0,0.0,0.0};
    pFeq[id].F0 = temp[15];
    pFeq[id].F0r = temp[16];
    pFeq[id].F0z = temp[17];
    return;
}

__global__ void rk2_kernel(fphsp* ph0a, fphsp* ph0b, fphsp* ph1) // phsp a = (a+b)/2
{
    int id = blockDim.x*blockIdx.x + threadIdx.x;
    ph0a[id].X = cmul(0.5, cadd(ph0a[id].X, ph0b[id].X));
    ph0a[id].vp = 0.5*(ph0a[id].vp + ph0b[id].vp);
    ph1[id].X = cadd(ph1[id].X, cmul(-1.0, ph0a[id].X));
    ph1[id].vp = ph1[id].vp - ph0a[id].vp;
}


// update particle location and velocity and check boundary
__global__ void updtkernel(fphsp* ph0, fphsp* ph1, feqdst* pf0eq, fpcrd* pcrd, 
fpcrd* pcrd_old, bool* pout, float* Aphi, bool updt_weight)
{
    int id = blockDim.x*blockIdx.x + threadIdx.x;
    //printf("part:%d, %d\n",id, pout[id]);
    fcylv crdnew = cxyz2cyl(cadd(ccyl2xyz(pcrd_old[id].crd), cmul(1.0,ph0[id].X)));
    bool box = crdnew.r < Rmin || crdnew.r > Rmax || crdnew.z < Zmin || crdnew.z > Zmax;
    if(box || !pout[id])
    {   
        pout[id] = false;
        return;
    }
    int rg = (int)((crdnew.r - Rmin)/dr);
    int zg = (int)((crdnew.z - Zmin)/dz);
    float ofr = (crdnew.r - (float)rg * dr - Rmin)/dr;
    float ofz = (crdnew.z - (float)zg * dz - Zmin)/dz;
    float coefrs[2] = {ofr, (float)1.0-ofr};
    float coefzs[2] = {ofz, (float)1.0-ofz};
    float psi = 0.0;
   
    for(int i = 0; i < 2; i++)
    for(int j = 0; j < 2; j++) {
        int point = Ny*(rg+i)+ zg+j;
        psi += Aphi[point]*coefrs[i]*coefzs[j];
    }
    if (psi/psiaxis < 0.1)
    {
        pout[id] = false;
        return;
    }
    pcrd[id].crd = crdnew;
    pcrd[id].vb = pcrd[id].vb + ph0[id].vp;
    if(updt_weight)
    {
        float h1r = crotate(ph1[id].X,pcrd_old[id].crd.p).x;
        pcrd[id].deltaw += ph1[id].X.z*pf0eq[id].F0z/pf0eq[id].F0;
        pcrd[id].deltaw += h1r*pf0eq[id].F0r/pf0eq[id].F0;
    }
    return;
}

// scatter to get a drift current field

__global__ void scatkernel(fpcrd* pc, float* phi_i, float* phi_q, int n)
{
    int idx = blockDim.x*blockIdx.x + threadIdx.x;
    fcylv vp = pc[idx].crd;

    int rg = (int)((vp.r - Rmin)/dr);
    int zg = (int)((vp.z - Zmin)/dz);
    float ofr = (vp.r - (float)rg * dr - Rmin)/dr;
    float ofz = (vp.z - (float)zg * dz - Zmin)/dz;
    
    float coefrs[4] = {ofr*(ofr-(float)1.0)*(ofr-(float)2.0)/-(float)6.0,
     (ofr+(float)1.0)*(ofr-(float)1.0)*(ofr-(float)2.0)/(float)2.0,
    (ofr+(float)1.0)*ofr*(ofr-(float)2.0)/-(float)2.0,
     (ofr+(float)1.0)*ofr*(ofr-(float)1.0)/(float)6.0};
    float coefzs[4] = {ofz*(ofz-(float)1.0)*(ofz-(float)2.0)/-(float)6.0,
     (ofz+(float)1.0)*(ofz-(float)1.0)*(ofz-(float)2.0)/(float)2.0,
     (ofz+(float)1.0)*ofz*(ofz-(float)2.0)/-(float)2.0,
     (ofz+(float)1.0)*ofz*(ofz-(float)1.0)/(float)6.0};
    
    for(int i = 0; i < 4; i++)
        for(int j = 0; j < 4; j++)
        {
            int point = 0;
            if(rg+i-1 < 0)
                point += Ny*(rg+i);
            else if(rg+i-1 > Nx)
                point += Ny*(rg+i-2);
            else
                point += Ny*(rg+i-1);
            if (zg+j-1 < 0)
                point += zg+j;
            else if (zg+j-1 > Nx)
                point += zg+j-2;
            else
                point += zg+j-1;
            phi_i[point] += coefrs[i]*coefzs[j]*pc[idx].deltaw*vp.r*cos(n*vp.p);
            phi_i[point] += coefrs[i]*coefzs[j]*pc[idx].deltaw*vp.r*sin(n*vp.p);
        }
}


void cupush_rk2(world wld, float dt, int nt, float omega, int n)
{
    dim3 block(1024);
    dim3 grid((wld.fastptcs->npt+block.x-1)/block.x);
    int nfpt = Nx*Ny;
    int npt = wld.fastptcs->npt;
    int nf = 18;
    float* flds = new float[nf*nfpt];
    float* pvals = new float[nf*npt];
    double** datapts = new double*[nf];

    /*  initialize fields */

    datapts[0] = wld.B0->getfr();
    datapts[1] = wld.B0->getfphi(); 
    datapts[2] = wld.B0->getfz();
    datapts[3] = wld.gB->getfr(); 
    datapts[4] = wld.gB->getfphi();
    datapts[5] = wld.gB->getfz(); 
    datapts[6] = wld.cb->getfr(); 
    datapts[7] = wld.cb->getfphi(); 
    datapts[8] = wld.cb->getfz();
    
    vfield gradF = wld.fast_distrib->gradient();
    vfiq E1 = vfiq(wld.phies->n);
    wld.phies->gradient(&E1);
    datapts[9] = E1.get_ifr();
    datapts[10] = E1.get_ifp();
    datapts[11] = E1.get_ifz();
    datapts[12] = E1.get_qfr();
    datapts[13] = E1.get_qfp();
    datapts[14] = E1.get_qfz();
    datapts[15] = wld.fast_distrib->getdata();
    datapts[16] = gradF.getfr();
    datapts[17] = gradF.getfz();

    for (int i = 0; i < nf; i++)
        for(int j = 0; j < nfpt; j++)
            flds[i*nfpt + j] = datapts[i][j];
    
    bool* pohost = new bool[nfpt];
    float* paphi = new float[nfpt];
    for(int j = 0; j < nfpt; j++)
        pohost[j] = false;

    double* pApsi = Apsi.getdata();
    float* apsi = new float[nfpt];
    for(int i = 0; i < nfpt; i++)
        apsi[i] = pApsi[i];
    /*initialize particles */
    fpcrd* pcrd = new fpcrd[wld.fastptcs->npt];
    fpcrd* pcrdnew = new fpcrd[wld.fastptcs->npt];
    bool* inside = new bool[wld.fastptcs->npt];
    for (int i = 0; i < wld.fastptcs->npt; i++)
    {
        pcrd[i].crd.r = wld.fastptcs->ppts[i].getcrd().r;
        pcrd[i].crd.p = wld.fastptcs->ppts[i].getcrd().p;
        pcrd[i].crd.z = wld.fastptcs->ppts[i].getcrd().z;
        pcrd[i].vb = wld.fastptcs->ppts[i].getvpara();
        pcrd[i].mu = wld.fastptcs->ppts[i].pmu();
        pcrd[i].deltaw = 0.0;
        inside[i] = true;
    }
    /* output density fields*/
    float* denf_i = new float[nfpt];
    float* denf_q = new float[nfpt];
    /* cuda variable allocate */

    float* cuflds, *cupvals; // cuda fields, interped values
    fpcrd* cupc; // cuda particle coordinate 
    fphsp* cuphsp_h0a, *cuphsp_h0b, *cuphsp_h1; //cuda phase space variables
    fpfvec* cupfs; // cuda ElectroMagnetic vector field variables
    feqdst* cupf0eq; // cuda F0 related values for particles 
   
    fpcrd* cupc_temp; // temporary coordinate for rk2
    bool* cuinside; // cuda bundary check
    float* cuaphi; // boundary check

    float* cudenf_i;
    float* cudenf_q;
    time_t start = time(NULL);
    cudaDeviceSynchronize();
    cudaMalloc((float**)&cuflds, sizeof(float)*nf*nfpt);
    cudaMalloc((float**)&cupvals, sizeof(float)*npt*nf);
    cudaMalloc((fpcrd**)&cupc, sizeof(fpcrd)*npt);
    cudaMalloc((fphsp**)&cuphsp_h0a, sizeof(fphsp)*npt);
    cudaMalloc((fphsp**)&cuphsp_h0b, sizeof(fphsp)*npt);
    cudaMalloc((fphsp**)&cuphsp_h1, sizeof(fphsp)*npt);
    cudaMalloc((fpfvec**)&cupfs, sizeof(fpfvec)*npt);
    cudaMalloc((feqdst**)&cupf0eq, sizeof(feqdst)*npt);
    cudaMalloc((bool**)&cuinside, sizeof(bool)*npt);
    cudaMalloc((float**)&cuaphi, sizeof(float)*nfpt);
    cudaMalloc((fpcrd**)&cupc_temp, sizeof(fpcrd)*npt);
    cudaMalloc((float**)&cudenf_i, sizeof(float)*nfpt);
    cudaMalloc((float**)&cudenf_q, sizeof(float)*nfpt);

    cudaMemcpy(cupc, pcrd, sizeof(fpcrd)*npt, cudaMemcpyHostToDevice);
    cudaMemcpy(cuflds, flds, sizeof(float)*nfpt*nf, cudaMemcpyHostToDevice);
    cudaMemcpy(cuinside, inside, sizeof(bool)*npt, cudaMemcpyHostToDevice);
    cudaMemcpy(cuaphi, apsi, sizeof(float)*nfpt, cudaMemcpyHostToDevice);
    cudaDeviceSynchronize();
    
    for(int i = 0; i < nt; i++)
    {
        intpkernel<<<grid, block>>>(cupc, cuflds, cupvals, 18, cuinside);       
        rogakernel<<<grid, block>>>(cupc, cupvals, cupfs, cupf0eq, false,n, omega, i*dt);    
        pushkernel<<<grid, block>>>(cupfs, cuphsp_h0a, cupc, cuinside, dt, 1.0,1.0);
        rogakernel<<<grid, block>>>(cupc, cupvals, cupfs, cupf0eq, true,n, omega, i*dt);    
        pushkernel<<<grid, block>>>(cupfs, cuphsp_h1, cupc, cuinside, dt, 1.0,1.0);        
        updtkernel<<<grid, block>>>(cuphsp_h0a, cuphsp_h1, cupf0eq, cupc_temp, cupc, cuinside ,cuaphi, false);    
        intpkernel<<<grid, block>>>(cupc_temp, cuflds, cupvals, 18, cuinside);  
        rogakernel<<<grid, block>>>(cupc_temp, cupvals, cupfs, cupf0eq, false,n, omega, i*dt);  
        pushkernel<<<grid, block>>>(cupfs, cuphsp_h0b, cupc, cuinside, dt, 1.0,1.0);
        rk2_kernel<<<grid, block>>>(cuphsp_h0a, cuphsp_h0b, cuphsp_h1);
        updtkernel<<<grid, block>>>(cuphsp_h0a, cuphsp_h1, cupf0eq, cupc, cupc, cuinside, cuaphi, true);
        if(i%100 == 0)
            scatkernel<<<grid, block>>>(cupc,cudenf_i,cudenf_q,n);
    }

    cudaMemcpy(pcrdnew, cupc, sizeof(fpcrd)*npt, cudaMemcpyDeviceToHost);
    cudaMemcpy(inside, cuinside, sizeof(bool)*npt, cudaMemcpyDeviceToHost);
    cudaMemcpy(denf_i,cudenf_i, sizeof(float)*nfpt, cudaMemcpyDeviceToHost);
    cudaMemcpy(denf_q, cudenf_q, sizeof(float)*nfpt, cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize();

    time_t stop = time(NULL);
    printf("-------------");
    printf("\nrunning on Nvidia RTX1660 super \ntime:%ds \nparticle number:%d \nsteps:%d\nnpt*nstp:%d\n",
    (int)difftime(stop,start),wld.fastptcs->npt,nt,wld.fastptcs->npt*nt);
    printf("-------------");

    for (int i = 0; i < wld.fastptcs->npt; i++)
    {
        double ri = pcrdnew[i].crd.r;
        double phii = pcrdnew[i].crd.p;
        double zi = pcrdnew[i].crd.z;

        wld.fastptcs->ppts[i].setcrd({ri,phii,zi},pcrdnew[i].vb);
        wld.fastptcs->wdelta[i] = pcrdnew[i].deltaw;
        wld.fastptcs->inside[i] = inside[i];
        if (i == 100)
            printf("\nfinal:%f,%f,%f,%f,%f\n", ri, phii,zi, pcrdnew[i].vb,pcrdnew[i].deltaw);
    }
    
    cudaFree(cuflds);
    cudaFree(cupvals);
    cudaFree(cupc);
    cudaFree(cuphsp_h0a);
    cudaFree(cuphsp_h0b);
    cudaFree(cuphsp_h1);
    cudaFree(cuflds);
    cudaFree(cupf0eq);
    cudaFree(cupc_temp);
    cudaFree(cuinside);
    cudaFree(cuaphi);
    cudaFree(cudenf_i);
    cudaFree(cudenf_q);
}


